import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { logger } from "./logger.js";
// 存储节点进度信息的映射
const nodeProgressMap = new Map();

// 创建进度条UI
function createProgressBar(node) {
    const container = document.createElement("div");
    container.style.padding = "10px";
    container.style.backgroundColor = "var(--comfy-input-bg)";
    container.style.borderRadius = "8px";
    container.style.marginTop = "10px";
    
    // 进度条容器
    const progressBarContainer = document.createElement("div");
    progressBarContainer.style.width = "100%";
    progressBarContainer.style.height = "20px";
    progressBarContainer.style.backgroundColor = "var(--comfy-input-bg)";
    progressBarContainer.style.border = "1px solid var(--border-color)";
    progressBarContainer.style.borderRadius = "4px";
    progressBarContainer.style.overflow = "hidden";
    
    // 进度条
    const progressBar = document.createElement("div");
    progressBar.style.width = "0%";
    progressBar.style.height = "100%";
    progressBar.style.backgroundColor = "var(--accent-color)";
    progressBar.style.transition = "width 0.3s ease";
    progressBarContainer.appendChild(progressBar);
    
    // 状态文本
    const statusText = document.createElement("div");
    statusText.style.marginTop = "5px";
    statusText.style.color = "var(--input-text)";
    statusText.textContent = "等待开始...";
    
    // 进度文本
    const progressText = document.createElement("div");
    progressText.style.marginTop = "5px";
    progressText.style.color = "var(--input-text)";
    progressText.textContent = "0 / 0 步";
    
    // Loss文本（如果有）
    const lossText = document.createElement("div");
    lossText.style.marginTop = "5px";
    lossText.style.color = "var(--input-text)";
    
    container.appendChild(progressBarContainer);
    container.appendChild(statusText);
    container.appendChild(progressText);
    container.appendChild(lossText);
    
    // 存储UI元素引用
    const nodeId = node.id !== undefined ? node.id : node.nodeId;
    if (!nodeId || nodeId === -1) {
        console.warn("[LoRA调试] 无效的节点ID, 将使用临时ID");
        // 使用时间戳作为临时ID
        node._tempId = node._tempId || Date.now();
        nodeProgressMap.set(node._tempId, {
            container,
            progressBar,
            statusText,
            progressText,
            lossText,
            isTemp: true
        });
        console.log("[LoRA调试] 使用临时ID存储进度条:", node._tempId);
    } else {
        nodeProgressMap.set(nodeId, {
            container,
            progressBar,
            statusText,
            progressText,
            lossText
        });
        console.log("[LoRA调试] 使用节点ID存储进度条:", nodeId);
    }
    
    return container;
}

// 更新进度条UI
function updateProgressBar(nodeId, progress) {
    console.log("[LoRA调试] 尝试更新进度条, nodeId:", nodeId);
    logger.analyzeObject(progress, "进度信息");
    
    // 尝试查找进度条
    let ui = nodeProgressMap.get(nodeId);
    if (!ui) {
        // 如果找不到，尝试遍历查找临时ID
        for (const [id, data] of nodeProgressMap.entries()) {
            if (data.isTemp) {
                ui = data;
                console.log("[LoRA调试] 找到临时进度条:", id);
                break;
            }
        }
    }
    
    if (!ui) {
        console.warn("[LoRA调试] 未找到节点进度条, nodeId:", nodeId);
        return;
    }
    
    const { current, total, status, loss } = progress;
    const percentage = (current / total * 100) || 0;
    
    ui.progressBar.style.width = `${percentage}%`;
    ui.statusText.textContent = status || "未知状态";
    ui.progressText.textContent = `${current} / ${total} 步`;
    
          if (loss !== undefined) {
          ui.lossText.textContent = `Loss: ${loss.toFixed(4)}`;
      }
  }
  
// 注册扩展
app.registerExtension({
    name: "Comfy.LoraTraining",
    
    async nodeCreated(node) {
        console.log("[LoRA调试] node对象:", node);
        console.log("[LoRA调试] node.comfyClass:", node.comfyClass);
        // 不做类型判断，所有节点都加小部件，便于调试
        try {
            // 添加标准slider进度条Widget
            const progressWidget = node.addWidget(
                "slider",
                "训练进度",
                50, // 默认50%
                null,
                {
                    min: 0,
                    max: 100,
                    step: 1,
                    readonly: true
                }
            );
            console.log("[LoRA调试] 进度条小部件已添加:", progressWidget);
            
            // 添加状态文本
            const statusWidget = node.addWidget("text", "训练状态", "", (value) => {
                console.log("[LoRA调试] 状态更新:", value);
            }, {
                multiline: false,
                readonly: true
            });
            console.log("[LoRA调试] 状态文本小部件已添加:", statusWidget);
            
            // 添加损失值文本
            const lossWidget = node.addWidget("text", "Loss", "", (value) => {
                console.log("[LoRA调试] Loss更新:", value);
            }, {
                multiline: false,
                readonly: true
            });
            console.log("[LoRA调试] Loss小部件已添加:", lossWidget);
            
            console.log("[LoRA调试] 所有组件已添加到节点");
            console.log("[LoRA调试] 当前节点widgets：", node.widgets);
        } catch (error) {
            console.error("[LoRA调试] 添加小部件时出错:", error);
        }
    }
    // getCustomWidgets 已注释，不再使用自定义Widget
}); 